Purpose

Fix a race on reused pinned CPU input buffers when engine steps overlap under the PP batch queue.

prepare_inputs_event (consumed by synchronize_input_prep()) was only created when async scheduling is enabled. But the PP batch queue creates exactly the same overlap (max_concurrent_batches == pp_size): input prep for step N+1 rewrites the reused pinned CPU tensors (seq_lens_cpu, query_start_loc_cpu, block-table staging, …) while step N's non_blocking H2D copies from those same buffers may not have executed yet. When the copy engine finally runs, step N's kernels consume step N+1's values.

For DSA sparse-MLA models this is fatal rather than silent: the indexer sizes its prefill-chunk buffers from step N's CPU values but its GPU kernels index with the (now newer, possibly larger) device seq_lens — out-of-bounds scatter → Triton Error [CUDA]: an illegal memory access was encountered. Reported in the field in #38476 (multi-node TP×PP, and single-node repro below). For dense models the same race can silently corrupt inputs.

The fix creates the event whenever max_concurrent_batches > 1, so synchronize_input_prep() guards every overlapped-step configuration (async scheduling ⇒ ≥ 2; PP ⇒ pp_size). Single-batch configurations are unchanged.

Why the symptom is so elusive

• CUDA_LAUNCH_BLOCKING=1 makes every H2D copy synchronous → pinned buffers are always consumed before reuse → 100% stable (verified). This masks the bug in exactly the debugging mode people reach for.
• TP-only has max_concurrent_batches == 1 → unaffected.
• The window is proportional to how far the GPU lags the CPU — huge on SM80 with Triton sparse-MLA prefill ([Attention] TRITON_MLA_SPARSE backend for SM80/SM121 sparse MLA (rebase & takeover of #38476) #47629), narrow-but-nonzero on H100.

Test Plan / Test Result

Repro hardware: 8×A800-80G PCIe (SM80), GLM-5.2 NVFP4 via #47629, TP=2 × PP=4, --no-async-scheduling, bf16 KV:

• Before: 9 concurrent long prefills (26k–86k tokens, 3-way concurrency) crash the engine within ~2 minutes (illegal memory access, surfacing at the indexer's build_prefill_chunk_metadata kernel launch). Reproduced twice; also reproduced by @Ph0enix89 on multi-node in [Feature] TRITON_MLA_SPARSE backend for SM8x/11x/12x DSA Sparse MLA Support #38476.
• After: same stress ×3 rounds — 27/27 requests complete, all needle-retrieval answers correct.
• Short-context PP serving, TP-only serving, and tool calling unaffected.

Notes

• Not a duplicate: no open issue/PR covers this (searched "pipeline parallel illegal memory access", "batch queue pinned race", prepare_inputs_event).
• The V2 model runner has its own input-prep path and is not touched here.
• AI assistance (Claude Code) was used to root-cause and draft this fix; the submitter reviewed every line and ran the tests above.

🤖 Generated with Claude Code